#!/usr/bin/env python3
"""
综合建筑安全检测模型训练脚本
支持多种安全违规行为检测的模型训练
"""

import os
import yaml
from ultralytics import YOLO
import torch
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path

def setup_environment():
    """设置训练环境"""
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    if torch.cuda.is_available():
        print(f"GPU名称: {torch.cuda.get_device_name(0)}")
        print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    
    return device

def create_comprehensive_dataset_structure():
    """创建综合数据集目录结构"""
    base_dirs = [
        'comprehensive_dataset/train/images',
        'comprehensive_dataset/train/labels',
        'comprehensive_dataset/valid/images', 
        'comprehensive_dataset/valid/labels',
        'comprehensive_dataset/test/images',
        'comprehensive_dataset/test/labels'
    ]
    
    for dir_path in base_dirs:
        os.makedirs(dir_path, exist_ok=True)
        print(f"创建目录: {dir_path}")
    
    print("\n数据集目录结构已创建。请按以下方式组织您的数据:")
    print("1. 将标注好的图片分别放入 train/images, valid/images, test/images")
    print("2. 将对应的标签文件放入 train/labels, valid/labels, test/labels")
    print("3. 标签文件格式: class_id x_center y_center width height (YOLO格式)")

def create_comprehensive_data_config():
    """创建综合安全检测的数据配置文件"""
    config = {
        'train': 'comprehensive_dataset/train/images',
        'val': 'comprehensive_dataset/valid/images', 
        'test': 'comprehensive_dataset/test/images',
        'nc': 12,  # 类别数量
        'names': [
            'Person',           # 0 - 人员
            'Helmet',           # 1 - 佩戴安全帽
            'No-Helmet',        # 2 - 未佩戴安全帽
            'Safety-Vest',      # 3 - 佩戴安全背心/工服
            'No-Safety-Vest',   # 4 - 未佩戴安全背心/工服
            'Smoking',          # 5 - 抽烟行为
            'Phone-Call',       # 6 - 接打电话
            'Fire',             # 7 - 火焰
            'Smoke',            # 8 - 烟雾
            'Danger-Zone',      # 9 - 危险区域
            'Safety-Equipment', # 10 - 安全设备
            'Vehicle'           # 11 - 工程车辆
        ]
    }
    
    with open('comprehensive_data.yaml', 'w', encoding='utf-8') as f:
        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    print("已创建综合安全检测数据配置文件: comprehensive_data.yaml")

def check_dataset():
    """检查数据集配置"""
    data_config = "comprehensive_data.yaml"
    
    if not os.path.exists(data_config):
        print(f"数据配置文件不存在，正在创建...")
        create_comprehensive_data_config()
    
    with open(data_config, 'r', encoding='utf-8') as f:
        data = yaml.safe_load(f)
    
    print("="*50)
    print("综合安全检测数据集信息:")
    print("="*50)
    print(f"类别数量: {data['nc']}")
    print(f"类别名称: {data['names']}")
    
    # 检查路径
    train_path = data['train']
    val_path = data['val']
    test_path = data['test']
    
    print(f"\n数据集路径:")
    print(f"训练集: {train_path}")
    print(f"验证集: {val_path}")
    print(f"测试集: {test_path}")
    
    # 检查数据是否存在
    for path_name, path in [('训练集', train_path), ('验证集', val_path), ('测试集', test_path)]:
        if os.path.exists(path):
            image_count = len([f for f in os.listdir(path) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))])
            print(f"{path_name}图片数量: {image_count}")
        else:
            print(f"⚠️ {path_name}路径不存在: {path}")
    
    return True

def train_comprehensive_model():
    """训练综合安全检测模型"""
    print("="*50)
    print("开始训练综合建筑安全检测模型")
    print("="*50)
    
    # 选择模型大小 - 根据任务复杂度选择
    model_options = {
        'n': 'yolov8n.pt',  # nano - 速度快，精度一般
        's': 'yolov8s.pt',  # small - 平衡
        'm': 'yolov8m.pt',  # medium - 精度较高
        'l': 'yolov8l.pt',  # large - 精度高，速度慢
        'x': 'yolov8x.pt'   # xlarge - 最高精度，最慢
    }
    
    # 由于类别较多，建议使用medium或large模型
    model_size = 'm'  # 可以根据需要调整
    model_name = model_options[model_size]
    
    print(f"使用模型: {model_name}")
    model = YOLO(model_name)
    
    # 训练配置 - 针对多类别检测优化
    training_config = {
        'data': 'comprehensive_data.yaml',
        'epochs': 150,              # 增加训练轮数，因为类别多
        'imgsz': 640,
        'batch': 16,                # 可根据GPU内存调整
        'device': 'auto',
        'workers': 4,
        'patience': 20,             # 增加耐心值
        'save': True,
        'save_period': 15,          # 更频繁地保存
        'cache': False,
        'optimizer': 'auto',
        
        # 学习率配置
        'lr0': 0.01,
        'lrf': 0.01,
        'momentum': 0.937,
        'weight_decay': 0.0005,
        'warmup_epochs': 5,         # 增加预热轮数
        'warmup_momentum': 0.8,
        'warmup_bias_lr': 0.1,
        
        # 数据增强 - 适合建筑场景
        'augment': True,
        'mosaic': 1.0,              # 启用马赛克增强
        'mixup': 0.1,               # 轻微混合增强
        'copy_paste': 0.1,          # 复制粘贴增强
        'degrees': 10.0,            # 旋转角度
        'translate': 0.1,           # 平移
        'scale': 0.5,               # 缩放
        'shear': 2.0,               # 剪切
        'perspective': 0.0001,      # 透视变换
        'flipud': 0.0,              # 垂直翻转（建筑场景不建议）
        'fliplr': 0.5,              # 水平翻转
        'hsv_h': 0.015,             # 色调变化
        'hsv_s': 0.7,               # 饱和度变化
        'hsv_v': 0.4,               # 亮度变化
        
        # 项目配置
        'project': 'runs/detect',
        'name': 'comprehensive_safety',
        'exist_ok': True,
        'verbose': True,
        
        # 类别平衡
        'cls': 0.5,                 # 分类损失权重
        'box': 7.5,                 # 边框损失权重
        'dfl': 1.5,                 # 分布焦点损失权重
    }
    
    print(f"训练配置:")
    for key, value in training_config.items():
        print(f"  {key}: {value}")
    
    # 开始训练
    print("\n开始训练...")
    results = model.train(**training_config)
    
    print("训练完成!")
    return results

def validate_model():
    """验证训练好的模型"""
    print("验证综合安全检测模型性能...")
    
    model_path = "runs/detect/comprehensive_safety/weights/best.pt"
    
    if not os.path.exists(model_path):
        print(f"错误: 找不到训练好的模型 {model_path}")
        return None
    
    model = YOLO(model_path)
    
    # 在验证集上评估
    metrics = model.val(data='comprehensive_data.yaml', verbose=True)
    
    print(f"\n模型性能指标:")
    print(f"mAP50: {metrics.box.map50:.3f}")
    print(f"mAP50-95: {metrics.box.map:.3f}")
    
    # 各类别性能
    if hasattr(metrics.box, 'ap_class_index'):
        print(f"\n各类别mAP50:")
        class_names = [
            'Person', 'Helmet', 'No-Helmet', 'Safety-Vest', 'No-Safety-Vest',
            'Smoking', 'Phone-Call', 'Fire', 'Smoke', 'Danger-Zone', 
            'Safety-Equipment', 'Vehicle'
        ]
        
        for i, class_name in enumerate(class_names):
            if i < len(metrics.box.ap50):
                print(f"  {class_name}: {metrics.box.ap50[i]:.3f}")
    
    return metrics

def analyze_training_results():
    """分析训练结果"""
    results_dir = "runs/detect/comprehensive_safety"
    
    if not os.path.exists(results_dir):
        print("未找到训练结果目录")
        return
    
    print(f"\n分析训练结果...")
    
    # 检查结果文件
    result_files = {
        'results.png': '训练曲线',
        'confusion_matrix.png': '混淆矩阵',
        'train_batch0.jpg': '训练批次示例',
        'val_batch0_pred.jpg': '验证预测示例',
        'val_batch0_labels.jpg': '验证标签示例'
    }
    
    existing_files = []
    for filename, description in result_files.items():
        filepath = os.path.join(results_dir, filename)
        if os.path.exists(filepath):
            existing_files.append((filepath, description))
            print(f"✓ {description}: {filename}")
        else:
            print(f"✗ {description}: {filename} (未找到)")
    
    # 显示训练结果图表
    if existing_files:
        plot_training_results(existing_files)

def plot_training_results(result_files):
    """绘制训练结果"""
    plt.style.use('default')
    
    n_files = len(result_files)
    if n_files == 0:
        return
    
    # 计算子图布局
    cols = min(3, n_files)
    rows = (n_files + cols - 1) // cols
    
    fig, axes = plt.subplots(rows, cols, figsize=(15, 5*rows))
    if rows == 1 and cols == 1:
        axes = [axes]
    elif rows == 1:
        axes = axes
    else:
        axes = axes.flatten()
    
    for i, (filepath, description) in enumerate(result_files):
        try:
            img = plt.imread(filepath)
            axes[i].imshow(img)
            axes[i].set_title(description, fontsize=12)
            axes[i].axis('off')
        except Exception as e:
            axes[i].text(0.5, 0.5, f'无法加载\n{description}', 
                        ha='center', va='center', transform=axes[i].transAxes)
            axes[i].axis('off')
    
    # 隐藏多余的子图
    for i in range(len(result_files), len(axes)):
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.savefig('comprehensive_training_summary.png', dpi=300, bbox_inches='tight')
    plt.show()

def generate_training_guide():
    """生成数据集准备指南"""
    guide = """
    
    ===============================================
    综合建筑安全检测数据集准备指南
    ===============================================
    
    1. 数据集类别说明:
       - Person (人员): 检测施工现场的工作人员
       - Helmet (佩戴安全帽): 正确佩戴安全帽的人员
       - No-Helmet (未佩戴安全帽): 未佩戴安全帽的人员头部
       - Safety-Vest (佩戴工服): 穿着反光背心/安全服的人员
       - No-Safety-Vest (未佩戴工服): 未穿安全服的人员
       - Smoking (抽烟): 吸烟行为
       - Phone-Call (接打电话): 使用手机通话的行为
       - Fire (火焰): 明火、火焰
       - Smoke (烟雾): 烟雾、粉尘
       - Danger-Zone (危险区域): 危险作业区域标识
       - Safety-Equipment (安全设备): 安全设备、警示标识
       - Vehicle (车辆): 工程车辆、机械设备
    
    2. 数据收集建议:
       - 每个类别建议至少收集500-1000张图片
       - 包含不同光照条件（早晨、中午、傍晚）
       - 包含不同天气条件（晴天、阴天、雨天）
       - 包含不同施工场景（高层建筑、道路施工、装修等）
       - 确保标注质量，边界框要紧贴目标对象
    
    3. 数据标注工具推荐:
       - LabelImg: 简单易用的标注工具
       - Roboflow: 在线标注平台，支持团队协作
       - CVAT: 功能强大的开源标注工具
    
    4. 数据增强策略:
       - 模型会自动进行数据增强
       - 包括旋转、缩放、颜色变换等
       - 针对建筑场景优化的增强参数
    
    5. 训练建议:
       - 建议使用GPU进行训练
       - 根据数据量调整训练轮数
       - 监控训练曲线，避免过拟合
       - 定期在验证集上评估性能
    
    ===============================================
    """
    
    print(guide)
    
    # 保存到文件
    with open('数据集准备指南.txt', 'w', encoding='utf-8') as f:
        f.write(guide)
    
    print("数据集准备指南已保存到: 数据集准备指南.txt")

def main():
    """主函数"""
    print("="*60)
    print("综合建筑安全检测模型训练系统")
    print("="*60)
    
    # 设置环境
    device = setup_environment()
    
    # 生成指南
    generate_training_guide()
    
    # 创建数据集结构
    create_comprehensive_dataset_structure()
    
    # 检查数据集
    if not check_dataset():
        return
    
    # 询问是否开始训练
    print("\n" + "="*50)
    user_input = input("是否开始训练? (y/n): ").lower().strip()
    
    if user_input == 'y' or user_input == 'yes':
        try:
            # 训练模型
            results = train_comprehensive_model()
            
            # 验证模型
            metrics = validate_model()
            
            # 分析结果
            analyze_training_results()
            
            print("\n" + "="*60)
            print("训练完成!")
            print("最佳模型保存在: runs/detect/comprehensive_safety/weights/best.pt")
            print("使用 comprehensive_safety_detector.py 进行检测")
            print("="*60)
            
        except Exception as e:
            print(f"训练过程中出现错误: {e}")
            print("请检查数据集路径和标注格式")
    
    else:
        print("请准备好数据集后再开始训练")
        print("参考数据集准备指南: 数据集准备指南.txt")

if __name__ == "__main__":
    main() 